refactor(trainer): extract the sync-loop template into BaseTrainer - #282
Closed
haonan3 wants to merge 1 commit into
Closed
refactor(trainer): extract the sync-loop template into BaseTrainer#282haonan3 wants to merge 1 commit into
haonan3 wants to merge 1 commit into
Conversation
2 tasks
Nine trainers each hand-copy the same outer loop (~700-900 duplicated lines repo-wide); the diffusion and AR bodies differ in exactly four lines. BaseTrainer now owns the loop as a template method (load ckpt → data fast-forward → wandb → baseline eval → per-rollout build/sync-gate/ train_step/log/eval/save → finally finish-wandb + on_train_end), with four hooks for the observed divergences: _wandb_extra, evaluate_baseline, evaluate_periodic, on_train_end (AR routes _shutdown_runtime through it, preserving its nested-finally ordering). DiffusionTrainer and ARTrainer drop their train() bodies (the template is line-equivalent, incl. the resume force-sync and the per-module progress logger name); PETrainer / UnifiedModelTrainer / SFTTrainer keep their own loops and migrate in follow-ups. Async/agentic loops are intentionally out of scope — their producer/consumer shapes are not this template.
haonan3
force-pushed
the
refactor/base-trainer-sync-template
branch
from
July 31, 2026 07:29
232b393 to
aac802f
Compare
2 tasks
Collaborator
Author
|
Closing by triage, not by verdict on the idea. The duplication is real and has grown since this was cut — ar / diffusion / pe / unified_model (and experimental refl) all hand-roll the same load→wandb→loop→save→finish skeleton. But the trainer surface is in an active evolution burst (AR PPO/GAE #259, value head #256, hi3 it2i #186), a template landed now adds hook-friction to every in-flight change, and this draft only migrated 2 of the 5 candidates with no GPU parity smoke. Re-derive when the AR-PPO wave settles, scoped to all five sync trainers in one pass — the branch ( |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Extracts the synchronous training loop into a
BaseTrainertemplate method. Nine trainers hand-copy the same outer skeleton today (~700–900 duplicated lines repo-wide, measured); the diffusion and ARtrain()bodies differ in exactly four lines.BaseTrainer.train()now owns the loop — load checkpoint → data fast-forward → wandb → baseline eval → per-rollout build / sync-gate /train_step/ log / periodic eval / save →finallyfinish-wandb +on_train_end— with four hooks for the observed divergences:_wandb_extra,evaluate_baseline,evaluate_periodic,on_train_end(AR routes_shutdown_runtimethrough it, preserving its nested-finallyordering).DiffusionTrainerandARTrainerdrop theirtrain()bodies (line-equivalent to the template, including the resume force-sync semantics and per-module progress-logger names).PETrainer/UnifiedModelTrainer/SFTTrainerkeep their own loops and migrate in follow-ups; async/agentic loops are intentionally out of scope (producer/consumer shapes, not this template). Independent of #210 (based on main).Test Plan
Static (CPU):
DiffusionTrainer.train is BaseTrainer.trainandARTrainer.train is BaseTrainer.trainassert-verified;ARTrainer.on_train_endoverride confirmed; compose-checkdiffusion/sd3/sd3_trainsideandar/qwen3_grpo_4b_base_dapo_sglangrc=0;pre-commit run --all-filesgreen.Pending GPU parity before merge:
eval_interval>0— baseline + periodic eval fire at the same steps; teardown runs on both clean exit and Ctrl-C.Compatibility / Risk
No config or entrypoint changes;
train(**kwargs)signature is unchanged for both migrated trainers. Behavioral risk is confined to the four hook seams — hence the explicit parity checklist above. Not-yet-migrated trainers are untouched.AI-assisted; reviewed and directed by the maintainer.